# -*- coding: utf-8 -*-
from __future__ import print_function

import tensorflow as tf
import numpy as np
import copy
import sys
import os
import json
from birnn_model import BirnnModel


def load_config():
    config = json.load(open("model_config.json"))
    return config


def load_vocab(path):
    word_to_id = json.load(open(path))
    id_to_word = {v: k for k, v in word_to_id.items()}
    return word_to_id, id_to_word


def main(_):
    prime = ["馒头"]
    vocab, id_to_char = load_vocab("word_to_id.json")
    config = load_config()
    ckpt_dir = "./models/"
    config['batch_size'] = 1

    # with tf.variable_scope('model'):
    model = BirnnModel(config, tf.contrib.learn.ModeKeys.INFER)
    checkpoint_file = tf.train.latest_checkpoint(ckpt_dir)

    print(checkpoint_file)
    sess = tf.Session()
    with sess.as_default():
        saver = tf.train.import_meta_graph("{}.meta".format(checkpoint_file))
        saver.restore(sess, checkpoint_file)
        tf.global_variables_initializer().run()
        res = model.sample(sess, id_to_char, vocab, 0, 10, prime)
        print(res)


if __name__ == '__main__':
    tf.app.run()
